
import os,sys,shutil
import pandas as pd
import tensorflow as tf
from PIL import Image
from xml_to_csv import xml_to_csv
from generate_tfrecord import  create_tf_example,split


data_set_root='/home/leo/Downloads/datas/road_car_split1'
data_set_root_csv='/home/leo/Downloads/datas/road_car_split1_csv'
data_set_root_record='/home/leo/Downloads/datas/road_car_split1_record'
label_file = 'label.txt'

if os.path.exists(data_set_root_csv) != True:
    os.makedirs(data_set_root_csv)

if os.path.exists(data_set_root_record) != True:
    os.makedirs(data_set_root_record)

for dir in os.listdir(data_set_root):
    gen_dir = os.path.join(data_set_root,dir)
    print("generating {}".format(gen_dir))
    xml_df = xml_to_csv(gen_dir)
    csv_file = os.path.join(data_set_root_csv,dir+".csv")
    xml_df.to_csv(os.path.join(csv_file), index=None)

    record_file = os.path.join(data_set_root_record,dir+".record")
    writer = tf.python_io.TFRecordWriter(record_file)
    examples = pd.read_csv(csv_file)
    grouped = split(examples, 'filename')
    for group in grouped:
        tf_example = create_tf_example(group, gen_dir)
        writer.write(tf_example.SerializeToString())
    writer.close()

shutil.copy(label_file,os.path.join(data_set_root_record,label_file))